#include <iostream>
#include <thread>
#include <vector>
#include <cstdlib>
#include <unistd.h> // for usleep
#include <chrono>   // for timing

#include "rpcExample/friend.pb.h"
#include "mprpcchannel.h"
#include "mprpccontroller.h"

// 每个线程的任务函数
void callRpcForUser(uint32_t userid, int count, const std::string& ip, short port, int& successCount, int& failureCount) {
    // 创建 stub 和 controller
    fixbug::FiendServiceRpc_Stub stub(new MprpcChannel(ip, port, true));
    MprpcController controller;

    // 创建请求和响应对象
    fixbug::GetFriendsListRequest request;
    fixbug::GetFriendsListResponse response;

    // 设置请求参数
    request.set_userid(userid);

    // 调用 RPC 方法 count 次
    while (count--) {
        // 发起 RPC 调用
        stub.GetFriendsList(&controller, &request, &response, nullptr);

        // 检查调用结果
        if (controller.Failed()) {
            ++failureCount; // 记录失败次数
        } else {
            if (0 == response.result().errcode()) {
                ++successCount; // 记录成功次数
            } else {
                ++failureCount; // 业务逻辑失败也算失败
            }
        }
    }
}

int main(int argc, char** argv) {
    std::string ip = "127.0.1.1";
    short port = 7788;

    int user_count = 500; // 用户数量
    int rpc_count = 1000; // 每个用户调用 RPC 的次数

    // 创建线程池
    std::vector<std::thread> threads;

    // 成功和失败计数
    int totalSuccessCount = 0;
    int totalFailureCount = 0;

    // 使用数组存储每个线程的成功和失败计数
    std::vector<int> successCounts(user_count, 0);
    std::vector<int> failureCounts(user_count, 0);

    // 记录开始时间
    auto startTime = std::chrono::steady_clock::now();

    // 为每个用户创建一个线程
    for (uint32_t i = 0; i < user_count; ++i) {
        uint32_t userid = 1000 + i;
        threads.emplace_back(callRpcForUser, userid, rpc_count, ip, port, std::ref(successCounts[i]), std::ref(failureCounts[i]));
    }

    // 等待所有线程完成
    for (auto& thread : threads) {
        thread.join();
    }

    // 记录结束时间
    auto endTime = std::chrono::steady_clock::now();

    // 统计总成功和失败次数
    for (int i = 0; i < user_count; ++i) {
        totalSuccessCount += successCounts[i];
        totalFailureCount += failureCounts[i];
    }

    // 计算总耗时（秒）
    std::chrono::duration<double> elapsedSeconds = endTime - startTime;

    // 计算 QPS
    double qps = totalSuccessCount / elapsedSeconds.count();

    // 输出结果
    std::cout << "压力测试完成！" << std::endl;
    std::cout << "总耗时: " << elapsedSeconds.count() << " 秒" << std::endl;
    std::cout << "总成功请求数: " << totalSuccessCount << std::endl;
    std::cout << "总失败请求数: " << totalFailureCount << std::endl;
    std::cout << "QPS: " << qps << " 请求/秒" << std::endl;

    return 0;
}